import time

import numpy as np
import pygame

from robots.piper import PiperArm


class XboxPiperController:
    def __init__(self, piper0: PiperArm, piper1: PiperArm):
        self.left_arm = piper0
        self.right_arm = piper1
        self.left_arm.ik.max_iter = 10
        self.right_arm.ik.max_iter = 10
        self._init_xbox_controller()

        # step length
        self.step_length = 0.01
        self.z_move_scale = 1
        self.deadzone = 0.5

        # temperal data
        self.fixed_left_z = None
        self.fixed_right_z = None
        self.left_xyz_limit = np.array([[0, -0.5, 0], [1, 1, 1]])
        self.right_xyz_limit = np.array([[0, -1, 0], [1, 0.5, 1]])

        # initialize gripper status
        self.left_arm.close_gripper()
        self.right_arm.close_gripper()
        self.left_arm_gripper_opened = False
        self.right_arm_gripper_opened = False

    def _init_xbox_controller(self):
        """初始化Xbox手柄"""
        pygame.init()
        pygame.joystick.init()

        if pygame.joystick.get_count() == 0:
            raise RuntimeError("未检测到Xbox手柄，请确保手柄已连接")

        self.joystick = pygame.joystick.Joystick(0)
        self.joystick.init()

        print(f"已连接Xbox手柄: {self.joystick.get_name()}")
        print(f"轴数量: {self.joystick.get_numaxes()}, 按钮数量: {self.joystick.get_numbuttons()}")
        print(f"帽子数量: {self.joystick.get_numhats()}")

        self.button_mapping = {
            0: "A",
            1: "B",
            2: "X",
            3: "Y",
            4: "L",
            5: "R",
            6: "Back",
            7: "Menu",
            8: "Xbox",
            9: "LDown",
            10: "RDown",
            11: "Share",
            # 4: "LB",  # 左肩键 - 关闭夹爪
            # 5: "RB",  # 右肩键 - 打开夹爪
            # 7: "Menu",  # 菜单键 - 重置位置
        }

        # 摇杆轴映射
        self.axis_mapping = {
            0: "Left Stick X",  # 左摇杆X轴 - 控制Y轴移动
            1: "Left Stick Y",  # 左摇杆Y轴 - 控制X轴移动
            2: "LT Trigger",  # 左触发器 - Z轴向下
            3: "Right Stick X",  # 右摇杆X轴 - 控制RZ轴旋转
            4: "Right Stick Y",  # 右摇杆Y轴 - 控制RY轴旋转
            5: "RT Trigger",  # 右触发器 - Z轴向上
        }

        self.hat_mapping = {
            (0, 1): "D-Pad Up",
            (0, -1): "D-Pad Down",
            (-1, 0): "D-Pad Left",
            (1, 0): "D-Pad Right",
            (-1, 1): "D-Pad Up-Left",
            (1, 1): "D-Pad Up-Right",
            (-1, -1): "D-Pad Down-Left",
            (1, -1): "D-Pad Down-Right",
        }

    def read_ctrl_from_xbox(self):

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    return False

        axes = [self.joystick.get_axis(i) for i in range(self.joystick.get_numaxes())]
        buttons = [self.joystick.get_button(i) for i in range(self.joystick.get_numbuttons())]
        hats = [self.joystick.get_hat(i) for i in range(self.joystick.get_numhats())]

        self.handle_joystick_control(axes, buttons, hats)

        return axes, buttons, hats

    def get_left_xyz(self):
        xyz = self.left_arm.get_effector_xeuler()[:3]
        if self.fixed_left_z is not None:
            xyz[2] = self.fixed_left_z
        return xyz

    def get_right_xyz(self):
        xyz = self.right_arm.get_effector_xeuler()[:3]
        if self.fixed_right_z is not None:
            xyz[2] = self.fixed_right_z
        return xyz

    def handle_joystick_control(self, axes, buttons, hats):
        """
        处理摇杆和D-pad控制 - V3.0优化版本

        Args:
            axes: 手柄轴数据列表
            hats: 手柄帽子数据列表
        """
        if len(axes) < 6:
            return

        triggered_left_arm = False
        triggered_right_arm = False

        # used for left arm
        left_stick_x = -axes[0]  # 左摇杆X轴（取反修正左右方向）
        left_stick_y = -axes[1]  # 左摇杆Y轴（反向）
        lt_trigger = axes[2]  # 左触发器

        # used for right arm
        right_stick_x = -axes[3]  # 右摇杆X轴（取反修正方向）
        right_stick_y = -axes[4]  # 右摇杆Y轴（修正方向）
        rt_trigger = axes[5]  # 右触发器

        xyz_left = self.get_left_xyz()
        xyz_right = self.get_right_xyz()

        # left arm
        if abs(left_stick_y) > self.deadzone:
            # move x
            xyz_left[0] += self.step_length * np.sign(left_stick_y)
            triggered_left_arm = True

        if abs(left_stick_x) > self.deadzone:
            # move y
            xyz_left[1] += self.step_length * np.sign(left_stick_x)
            triggered_left_arm = True

        if self.fixed_left_z is None:
            if buttons[4] == 1:
                # move z up
                xyz_left[2] += self.step_length * self.z_move_scale
                triggered_left_arm = True

            if lt_trigger > self.deadzone:
                # move z down
                xyz_left[2] -= self.step_length * self.z_move_scale
                triggered_left_arm = True

        if buttons[9] == 1:
            # fixed left z
            print("fixed left z")
            if self.fixed_left_z is None:
                self.fixed_left_z = xyz_left[2]
                print(f"Left z fixed to {self.fixed_left_z}")
            else:
                self.fixed_left_z = None
                print(f"Left z released")

        if buttons[6] == 1:
            if self.left_arm_gripper_opened:
                self.left_arm.close_gripper()
                self.left_arm_gripper_opened = False
            else:
                self.left_arm.open_gripper()
                self.left_arm_gripper_opened = True

        """
        right arm
        """

        if abs(right_stick_y) > self.deadzone:
            # move x
            xyz_right[0] += self.step_length * np.sign(right_stick_y)
            triggered_right_arm = True

        if abs(right_stick_x) > self.deadzone:
            # move y
            xyz_right[1] += self.step_length * np.sign(right_stick_x)
            triggered_right_arm = True

        if self.fixed_right_z is None:
            if buttons[5] == 1:
                # move z up
                xyz_right[2] += self.step_length * self.z_move_scale
                triggered_right_arm = True

            if rt_trigger > self.deadzone:
                # move z down
                xyz_right[2] -= self.step_length * self.z_move_scale
                triggered_right_arm = True

        if buttons[10] == 1:
            # fixed right z
            if self.fixed_right_z is None:
                self.fixed_right_z = xyz_right[2]
                print(f"Right z fixed to {self.fixed_right_z}")
            else:
                self.fixed_right_z = None
                print(f"Right z released")
        if buttons[7] == 1:
            if self.right_arm_gripper_opened:
                self.right_arm.close_gripper()
                self.right_arm_gripper_opened = False
            else:
                self.right_arm.open_gripper()
                self.right_arm_gripper_opened = True

        if triggered_left_arm:
            xyz_left = np.clip(xyz_left, self.left_xyz_limit[0], self.left_xyz_limit[1])
            self.left_arm.move_to(xyz_left)
            print(f"Move left armto {xyz_left}", end="\r")

        if triggered_right_arm:
            xyz_right = np.clip(xyz_right, self.right_xyz_limit[0], self.right_xyz_limit[1])
            self.right_arm.move_to(xyz_right)
            print(f"Move right armto {xyz_right}", end="\r")

        # go zero
        if buttons[11] == 1:
            self.left_arm.go_zero()
            self.right_arm.go_zero()
            self.left_arm.close_gripper()
            self.right_arm.close_gripper()
            self.left_arm_gripper_opened = False
            self.right_arm_gripper_opened = False
            self.fixed_left_z = None
            self.fixed_right_z = None


if __name__ == "__main__":
    self = XboxPiperController(PiperArm("can0", wait_execution=False), PiperArm("can1", wait_execution=False))
    while True:
        axes, buttons, hats = self.read_ctrl_from_xbox()
        # print(axes, buttons, hats)
        time.sleep(0.01)
